# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

import os

from absl import app, flags

import pandas
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

FLAGS = flags.FLAGS


def plot_premium(df, exchange, csv_root, product):
  plt.figure(figsize=(20, 10))
  plt.title('Fcoin -> %s %s premium' % (exchange, product))
  df['bid_ask_premium_x_y'] = df['true_bid_x'] - df['true_ask_y']
  df['bid_ask_premium_y_x'] = df['true_bid_y'] - df['true_ask_x']
  df['mid_premium'] = (df['true_ask_x'] + df['true_bid_x'] - df['true_ask_y']
                       - df['true_bid_y']) / 2

  bid_ask_premium_x_y, = plt.plot(df['bid_ask_premium_x_y'], color='blue')
  bid_ask_premium_y_x, = plt.plot(df['bid_ask_premium_y_x'], color='green')
  # mid_premium, = plt.plot(df['mid_premium'], color='red')

  plt.legend([bid_ask_premium_x_y, bid_ask_premium_y_x],
             ['Fcoin_true_bid - %s_true_ask' % exchange, '%s_true_bid - Fcoin_true_ask' % exchange])
  plt.ylabel('premium USDT')
  plt.grid(True)
  plt.savefig(os.path.join(csv_root, 'Fcoin_%s_%s_premium.png' % (exchange, product)))


def merge_dfs(df1, df2):
  return pandas.merge(df1, df2, how='inner', left_index=True, right_index=True)


def sample_df_per_second(df):
  return df.resample('S', closed='left', label='left').last()


def read_csv_into_df(csv_root, csv):
  csv_dir = os.path.join(csv_root, csv)
  df = pandas.read_csv(csv_dir, sep=',', header=0)
  df['datetime'] = pandas.to_datetime(df['timestamp'], unit='ns')
  df = df.set_index('datetime')
  return df


def main(argv):
  csv_root = FLAGS.csv_root
  assert csv_root, '--csv_root must be specified.'

  products = ['BTC-USDT', 'ETC-USDT', 'ETH-USDT']
  for product in products:
    fcoin_price_df = read_csv_into_df(csv_root, '%s.%s.csv' % (product, 'Fcoin'))
    for exchange in ['Binance', 'Bitfinex', 'Huobi', 'Okex']:
      other_price_df = read_csv_into_df(csv_root, '%s.%s.csv' % (product, exchange))
      sampled_df1 = sample_df_per_second(fcoin_price_df)
      sampled_df2 = sample_df_per_second(other_price_df)
      merged_df = merge_dfs(sampled_df1, sampled_df2)
      plot_premium(merged_df, exchange, csv_root, product)

  return 0


if __name__ == '__main__':
  flags.DEFINE_string('csv_root', None, 'Input csv files root directory.')

  app.run(main)
